package com.dolphin.aiagentplatform.ai.client;

import com.dolphin.aiagentplatform.ai.constant.AgentTypeConst;
import com.dolphin.aiagentplatform.ai.constant.ModelTypeConst;
import com.dolphin.aiagentplatform.ai.repository.ChatIdRepository;
import jakarta.annotation.Resource;
import lombok.Data;
import org.springframework.ai.chat.client.ChatClient;
import org.springframework.ai.chat.client.advisor.MessageChatMemoryAdvisor;
import org.springframework.ai.chat.client.advisor.SimpleLoggerAdvisor;
import org.springframework.ai.chat.memory.ChatMemory;
import org.springframework.ai.openai.OpenAiChatModel;
import org.springframework.ai.openai.OpenAiChatOptions;
import org.springframework.ai.openai.api.OpenAiApi;
import org.springframework.util.StringUtils;
import reactor.core.publisher.Flux;

import static org.springframework.ai.chat.client.advisor.AbstractChatMemoryAdvisor.CHAT_MEMORY_CONVERSATION_ID_KEY;

@Data
public abstract class AbstractChatClientGenerator implements ChatClientStrategy {

    @Resource
    protected ChatMemory chatMemory;

    @Resource
    protected ChatIdRepository inMemoryChatIdRepository;

    protected abstract String getBaseUrl();

    protected abstract String getApiKey();

    protected abstract String getModelName();

    public abstract AgentTypeConst getChatType();

    public abstract ModelTypeConst getModelType();

    protected ChatClient generateClient(Double temperature, Integer maxTokens) {
        OpenAiApi api = new OpenAiApi(getBaseUrl(), getApiKey());
        OpenAiChatOptions options = OpenAiChatOptions.builder()
                .temperature(temperature)
                .maxTokens(maxTokens)
                .model(getModelName())
                .build();
        OpenAiChatModel openAiChatModel = new OpenAiChatModel(api, options);
        return ChatClient.builder(openAiChatModel)
                .defaultOptions(options)
                .defaultAdvisors(
                        new SimpleLoggerAdvisor(),
                        new MessageChatMemoryAdvisor(chatMemory))
                .build();
    }

    @Override
    public Flux<String> chat(String systemPrompt, String userPrompt,
                             Double temperature, Integer maxTokens, String chatId) {
        ChatClient client = generateClient(temperature, maxTokens);
        inMemoryChatIdRepository.save(getChatType().toString(), chatId);

        // 创建提示构建器
        var promptBuilder = client.prompt();

        // 只有当systemPrompt有内容时才添加system消息
        if (StringUtils.hasText(systemPrompt)) {
            promptBuilder.system(systemPrompt);
        }

        // 添加用户消息并返回流
        return promptBuilder
                .user(userPrompt)
                .advisors(advisorSpec -> advisorSpec.param(CHAT_MEMORY_CONVERSATION_ID_KEY, chatId))
                .stream()
                .content();
    }
} 